#include "Searcher.h"

#include <clang/AST/RecursiveASTVisitor.h>

#include "Coder.h"

using namespace std;
using namespace clang;
using namespace rocket;

/**
 * @brief The RecordSearcher class 根据名检索类定义
 */
class RecordSearcher : public RecursiveASTVisitor<RecordSearcher> {
 protected:
  string rname;
  common::CXXRecordOpts opts;
  RecordDecl* rd = nullptr;

 public:
  RecordSearcher(string _rname) : rname(std::move(_rname)) {}
  RecordSearcher(string _rname, common::CXXRecordOpts _opts)
      : rname(std::move(_rname)), opts(std::move(_opts)) {}

 public:
  RecordDecl* getTargetRecordDecl() const { return rd; }
  CXXRecordDecl* getTargetCXXRecordDecl() const {
    return dyn_cast<CXXRecordDecl>(rd);
  }

 public:
  bool VisitRecordDecl(clang::RecordDecl* rd) {
    if (rname == common::getFullNDeclName(rd)) {
      this->rd = rd->getDefinition();
      return false;
    }
    return true;
  }

  bool VisitCXXRecordDecl(clang::CXXRecordDecl* cxxrd) {
    if (rname == common::getFullNDeclName(cxxrd)) {
      this->rd = cxxrd->getDefinition();
      return false;
    }
    return true;
  }
};

/**
 * @brief The FuncSearcher class
 * 根据函数mangled name或函数名及其形参类型检索函数定义
 */
class FuncSearcher : public RecursiveASTVisitor<FuncSearcher> {
 protected:
  string fname;
  vector<string> fargts;
  FunctionDecl* fd = nullptr;

  enum class SEARCH_METHOD {
    MANGLED,
    NAMEARGS,
  } skind;

 private:
  bool funcMatch(FunctionDecl* fd) {
    switch (skind) {
      case SEARCH_METHOD::MANGLED: {
        if (fname == common::getFullNDeclName(fd, true)) {
          this->fd = fd;
          return true;
        }
        break;
      }
      case SEARCH_METHOD::NAMEARGS: {
        if (fname == common::getFullNDeclName(fd) &&
            (fd->param_size() == fargts.size() ||
             fd->param_size() == (fargts.size() - 1))) {
          unsigned int i = 0;
          for (; i < fd->param_size(); ++i) {
            auto pd = fd->getParamDecl(i);

            const string& v_argn = fargts[i];
            const string& f_argn =
                common::getFullQTypeName(pd->getType(), fd->getASTContext());
            if (v_argn != f_argn) return false;
          }

          if (fd->isVariadic()) {
            return fargts[i] == "...";
          } else if (fd->param_size() < fargts.size()) {
            return false;
          }
          return true;
        }
        break;
      }
    }
    return false;
  }

 public:
  FuncSearcher(string _fname)
      : fname(std::move(_fname)), skind(SEARCH_METHOD::MANGLED) {}
  FuncSearcher(string _fname, vector<string> _argts)
      : fname(std::move(_fname)),
        fargts(std::move(_argts)),
        skind(SEARCH_METHOD::NAMEARGS) {}

 public:
  FunctionDecl* getTargetFuncDecl() const { return fd; }

 public:
  bool VisitFunctionDecl(FunctionDecl* fd) {
    if (funcMatch(fd)) {
      this->fd = fd;
      return false;
    }
    return true;
  }

  bool VisitFunctionTemplateDecl(FunctionTemplateDecl* ftd) {
    auto fd = ftd->getTemplatedDecl();
    if (funcMatch(fd)) {
      this->fd = fd;
      return false;
    }
    return true;
  }

  bool VisitClassTemplateDecl(ClassTemplateDecl* ctd) {
    for (auto spec : ctd->specializations()) {
      for (auto spec_sub_decl : spec->decls()) {
        if (auto method = dyn_cast<FunctionDecl>(spec_sub_decl)) {
          this->VisitFunctionDecl(method);
        } else if (auto tmethod =
                       dyn_cast<FunctionTemplateDecl>(spec_sub_decl)) {
          this->VisitFunctionTemplateDecl(tmethod);
        }
      }
    }
    return true;
  }
};

/**
 * @brief The FuncsSearcher class
 * 根据函数名检索同名函数
 */
class FuncsSearcher : public RecursiveASTVisitor<FuncSearcher> {
 protected:
  string fname;
  set<FunctionDecl*> fds = {};

 public:
  FuncsSearcher(string _fname) : fname(_fname) {}

 public:
  set<FunctionDecl*> getTargetFuncDecls() const { return fds; }

 public:
  bool VisitFunctionDecl(FunctionDecl* fd) {
    if (common::getFullNDeclName(fd) == fname) {
      fds.insert(fd);
    }
    return true;
  }
};

/* ======================================================================= */

RecordDecl* rocket::common::findNamedRecord(Decl* decl, string name) {
  if (!decl) return nullptr;
  RecordSearcher rs(name);
  rs.TraverseDecl(decl);
  return rs.getTargetRecordDecl();
}

CXXRecordDecl* common::findNamedCXXRecord(Decl* decl, string name,
                                          common::CXXRecordOpts opts) {
  if (!decl) return nullptr;
  RecordSearcher rs(name, opts);
  rs.TraverseDecl(decl);
  return rs.getTargetCXXRecordDecl();
}

FunctionDecl* common::findMangledFunction(Decl* decl, string mangled) {
  if (!decl) return nullptr;
  FuncSearcher fs(mangled);
  fs.TraverseDecl(decl);
  return fs.getTargetFuncDecl();
}

std::set<FunctionDecl*> common::findNamedFunctions(Decl* decl, string fname) {
  if (!decl) return {};
  FuncsSearcher fss(fname);
  fss.TraverseDecl(decl);
  return fss.getTargetFuncDecls();
}

FunctionDecl* common::findNamedFunction(Decl* decl, string fname,
                                        std::vector<string> fargts) {
  if (!decl) return nullptr;
  FuncSearcher fs(fname, fargts);
  fs.TraverseDecl(decl);
  return fs.getTargetFuncDecl();
}

const CXXRecordDecl* common::findRecordNamedBase(const CXXRecordDecl* rd,
                                                 string name) {
  if (!rd) return nullptr;
  if (common::getFullNDeclName(rd, false) == name) return rd;
  for (auto base : rd->bases()) {
    auto base_rd = base.getType()->getAsCXXRecordDecl();
    assert(base_rd);
    auto based = findRecordNamedBase(base_rd, name);
    if (based) return based;
  }
  return nullptr;
}

void common::getAllRecordBases(const CXXRecordDecl* rd,
                               std::set<const CXXRecordDecl*>& bases) {
  if (!rd || bases.count(rd) > 0) return;
  bases.insert(rd);
  if (!rd->hasDefinition()) return;
  for (auto base_spec : rd->bases()) {
    auto base_type = base_spec.getType();
    if (base_type->isRecordType()) {
      common::getAllRecordBases(base_type->getAsCXXRecordDecl(), bases);
    }
  }
}

std::set<const CXXRecordDecl*> common::getAllRecordBases(
    const CXXRecordDecl* rd) {
  if (!rd) return {};
  if (!rd->hasDefinition()) return {};
  std::set<const CXXRecordDecl*> bases;
  getAllRecordBases(rd, bases);
  return bases;
}

static void getAllRecordBasesWideFirstList(
    const CXXRecordDecl* rd, std::vector<const CXXRecordDecl*>& list,
    std::set<const CXXRecordDecl*>& set) {
  if (!rd) return;
  if (!rd->hasDefinition()) return;
  std::vector<const CXXRecordDecl*> cxxrd_buffer;
  for (auto base_spec : rd->bases()) {
    auto base_type = base_spec.getType();
    if (base_type->isRecordType()) {
      auto cxxrd = base_type->getAsCXXRecordDecl();
      if (set.count(cxxrd) == 0) {
        set.insert(cxxrd);
        list.push_back(cxxrd);
        cxxrd_buffer.push_back(cxxrd);
      }
    }
  };
  for (auto base_rd : cxxrd_buffer) {
    getAllRecordBasesWideFirstList(base_rd, list, set);
  }
}

std::vector<const CXXRecordDecl*> common::getAllRecordBasesList(
    const CXXRecordDecl* rd) {
  if (!rd) return {};
  if (!rd->hasDefinition()) return {};
  std::vector<const CXXRecordDecl*> list;
  std::set<const CXXRecordDecl*> set;
  list.push_back(rd);
  set.insert(rd);
  getAllRecordBasesWideFirstList(rd, list, set);
  return list;
}

std::map<const Type*, std::set<const Expr*> > common::getAllFunctionExceptions(
    const FunctionDecl* fd) {
  fd = fd ? fd->getDefinition() : fd;
  if (!fd) return {};

  class ExceptionVisitor : public clang::RecursiveASTVisitor<ExceptionVisitor> {
   public:
    std::map<const Type*, std::set<const Expr*> >& exceptions;

    ExceptionVisitor(std::map<const Type*, std::set<const Expr*> >& exceptions)
        : exceptions(exceptions) {}

    bool VisitCXXThrowExpr(clang::CXXThrowExpr* cxxte) {
      auto te_expr = cxxte->getSubExpr();
      if (te_expr) {
        exceptions[te_expr->getType().getTypePtr()].insert(te_expr);
      }
      return true;
    }
  };

  std::map<const Type*, std::set<const Expr*> > exceptions;
  ExceptionVisitor visitor(exceptions);
  visitor.TraverseStmt(fd->getBody());
  return exceptions;
}
